# -*- coding: utf-8 -*-
"""
Created on Mon Oct  8 14:32:52 2018

@author: luolei

单个样本在72小时上的预测效果和与真实值的对比
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import sys

sys.path.append('../..')

from lib.eval_model import *
from lib.eval_model.nn_evaluation import model_prediction

sample_num = 3800


#%% 主程序
if __name__ == '__main__':
	# 参数设置
	nn_model = load_models()
	
	# 构造测试数据
	X_test, y_test, total_cols_n = build_test_samples_and_targets()
	
	# 提取单条样本并进行模型预测
	x_test, y_test = X_test[sample_num - 1: sample_num, :], y_test[sample_num - 1: sample_num, :]
	y_real = y_test[0, :].reshape(1, -1)
	y_pred = model_prediction(x_test, nn_model)

	# 各污染物数据切分
	y_real_dict, y_pred_dict = {}, {}
	for i in range(len(target_cols)):
		y_real_dict[target_cols[i]] = y_real[:, i * pred_dim: (i + 1) * pred_dim]
		y_pred_dict[target_cols[i]] = y_pred[:, i * pred_dim: (i + 1) * pred_dim]

	# 还原为真实值
	for col in target_cols:
		bounds = var_bounds[col]
		y_real_dict[col] = y_real_dict[col] * (bounds[1] - bounds[0]) + bounds[0]
		y_pred_dict[col] = y_pred_dict[col] * (bounds[1] - bounds[0]) + bounds[0]

	# 载入计算得到的各时间预测步loss数据
	widths = np.linspace(0.5, 8.0, 5)  # **与alphas对应，长度一致
	eval_records = pd.read_csv('../../file/model/nn_evaluation_results.csv')
	loss_dict = {}
	for col in target_cols:
		loss_dict[col] = {}
		for width in widths:
			loss_dict[col][width] = width * np.array(eval_records.loc[:, col + '_mae']).flatten()

	# 计算预测结果以及上下界数据
	pred_curves = {}
	for col in target_cols:
		pred_curves[col] = {}
		pred_curves[col]['middle'] = y_pred_dict[col].flatten()
		
		pred_curves[col]['upper'] = {}
		pred_curves[col]['lower'] = {}
		for width in widths:
			pred_curves[col]['upper'][width] = pred_curves[col]['middle'] + loss_dict[col][width]
			pred_curves[col]['lower'][width] = pred_curves[col]['middle'] - loss_dict[col][width]
			lower_bound = pred_curves[col]['lower'][width]
			lower_bound[lower_bound < 0] = 0
	
	# 显示该时刻的预测效果
	alphas = np.logspace(1, 11, 5, base = 0.7)  # **与widths对应，长度一致
	plt.figure('pred results', figsize = [4, 2 * len(target_cols)])
	for col in target_cols:
		plt.subplot(len(target_cols), 1, target_cols.index(col) + 1)
		# plt.plot(pred_curves[col]['middle'], 'b', label = 'pred')           # 预测均值
		# plt.plot(pred_curves[col]['upper'], 'b--', linewidth = 0.5)
		# plt.plot(pred_curves[col]['lower'], 'b--', linewidth = 0.5)
		
		x = np.arange(len(pred_curves[col]['middle']))      # 绘制偏差
		for i in range(len(widths)):
			width = widths[i]
			plt.fill_between(
				x,
				pred_curves[col]['lower'][width],
				pred_curves[col]['upper'][width],
				facecolor = 'cornflowerblue',
				alpha = alphas[i]
			)
			
		plt.plot(y_real_dict[col].flatten(), c = 'red', linewidth = 0.8, label = 'real')           # 真实值
		plt.ylabel(col, fontsize = 10)
		plt.xticks(fontsize = 6)
		plt.yticks(fontsize = 6)
		plt.legend(fontsize = 6, loc = 'upper right')
	plt.xlabel('pred time step', fontsize = 10)
	plt.tight_layout()
	plt.savefig('../../graph/pred_effect.png', dpi = 450)
